""" 优化版OCR web服务 """
import json
import os
import tempfile
from pathlib import Path
from threading import Thread
from time import time
from concurrent.futures import as_completed,ThreadPoolExecutor
from typing import Literal, Dict, List
from functools import partial
from dataclasses import dataclass
import psutil

from bottle import route, run, post, request, BaseRequest, response

from PPOCR_api import GetOcrApi
from RapidOCR_api import OcrAPI

ALLOWED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.pdf'}
IMAGE_EXTENSIONS = {'.png', '.jpg', '.jpeg'}

# OCR配置
SERVER_DIR = Path(__file__).resolve().parent
Rapid_OCR_PATH = str(SERVER_DIR.joinpath("OCRCore/RapidOCR-json_v0.2.0/RapidOCR-json.exe"))
Paddle_OCR_PATH = str(SERVER_DIR.joinpath(
    "OCRCore/PaddleOCR-json_v1.4.1_windows_x64/PaddleOCR-json_v1.4.1/PaddleOCR-json.exe"))
TEST_OCR_FILE_PATH = str(SERVER_DIR.joinpath("images/test_ocr.png"))

# 性能优化配置
MAX_WORKERS = min(psutil.cpu_count(), 8)
SUPPORTED_OCR_TYPES = []  # 存储可用的OCR引擎类型
OCR_PROCESS_POOL: Dict[str, List] = {}  # 动态初始化进程池
PDF_CONVERT_WORKERS = 2


@dataclass
class OCRTask:
    type: str
    filename: str
    data: bytes or str = None
    ocr_type: str = "RapidOCR"
    original_filename: str = None
    page_num: int = None
    error: str = None

def detect_supported_ocr():
    """
    检测支持的OCR引擎
    :return:
    """
    global SUPPORTED_OCR_TYPES,OCR_PROCESS_POOL
    SUPPORTED_OCR_TYPES = []
    supported_ocr = []
    for ocr_core_type in [ "RapidOCR","PaddleOCR"]:
        try:
            if ocr_core_type == "RapidOCR":
                ocr = OcrAPI(Rapid_OCR_PATH)
                result = ocr.run(TEST_OCR_FILE_PATH)
            else:
                ocr = GetOcrApi(Paddle_OCR_PATH)
                result = ocr.run(TEST_OCR_FILE_PATH)
            if result:
                OCR_PROCESS_POOL[ocr_core_type] = []
                supported_ocr.append(ocr_core_type)
        except Exception as e:
            print(f"{ocr_core_type} 检测失败: {str(e)}")
    SUPPORTED_OCR_TYPES = supported_ocr


@route('/ocr_test')
def ocr_test():
    """
    测试ocr是否可用并初始化进程池
    :return: 返回可用引擎列表和初始化状态
    """
    try:
        # 先清空原有支持列表和进程池
        global SUPPORTED_OCR_TYPES, OCR_PROCESS_POOL
        SUPPORTED_OCR_TYPES = []
        OCR_PROCESS_POOL = {}

        # 检测可用引擎
        supported_ocr = []
        for ocr_core_type in ["PaddleOCR", "RapidOCR"]:
            try:
                if ocr_core_type == "RapidOCR":
                    ocr = OcrAPI(Rapid_OCR_PATH)
                    result = ocr.run(TEST_OCR_FILE_PATH)
                else:
                    ocr = GetOcrApi(Paddle_OCR_PATH)
                    result = ocr.run(TEST_OCR_FILE_PATH)
            except Exception as e:
                print(f"{ocr_core_type} 检测失败: {str(e)}")
            finally:
                if result and result["code"] == 100:  # 假设100是成功代码
                    supported_ocr.append(ocr_core_type)
                    OCR_PROCESS_POOL[ocr_core_type] = []  # 初始化该引擎的进程池
                if ocr:
                    ocr.stop() if hasattr(ocr, 'stop') else ocr.exit()

        SUPPORTED_OCR_TYPES = supported_ocr

        # 如果有可用引擎则初始化进程池
        if SUPPORTED_OCR_TYPES:
            # 开启单个新线程，用于初始化OCR进程池
            # with ThreadPoolExecutor(max_workers=1) as executor:
            #     executor.submit(init_ocr_pool)
            t = Thread(target=init_ocr_pool, daemon=True)
            t.start()
            return json.dumps({
                "code": 200,
                "msg": "OCR可用并已初始化进程池",
                "data": {
                    "supported_ocr": SUPPORTED_OCR_TYPES,
                    "initialized": True
                }
            })
        else:
            return json.dumps({
                "code": 404,
                "msg": "没有可用的OCR引擎",
                "data": {
                    "supported_ocr": [],
                    "initialized": False
                }
            })

    except Exception as e:
        return json.dumps({
            "code": 500,
            "msg": f"检测过程中发生错误: {str(e)}",
            "data": {
                "supported_ocr": [],
                "initialized": False
            }
        })




def init_ocr_pool():
    """预初始化OCR进程池"""

    for ocr_type in SUPPORTED_OCR_TYPES:
        OCR_PROCESS_POOL[ocr_type] = []
        for _ in range(MAX_WORKERS):
            if ocr_type == "RapidOCR":
                OCR_PROCESS_POOL[ocr_type].append(OcrAPI(Rapid_OCR_PATH))
            else:
                OCR_PROCESS_POOL[ocr_type].append(GetOcrApi(Paddle_OCR_PATH))


def get_ocr_instance(ocr_type: str):
    """获取OCR实例"""

    if not OCR_PROCESS_POOL.get(ocr_type):
        init_ocr_pool()

    return OCR_PROCESS_POOL[ocr_type].pop(0)


def release_ocr_instance(ocr_type: str, instance):
    """释放OCR实例"""
    OCR_PROCESS_POOL[ocr_type].append(instance)


def process_image_memory_task(file_bytes, ocr_core_type: str = "RapidOCR"):
    """优化版内存图片处理"""
    ocr_core = None
    try:
        ocr_core = get_ocr_instance(ocr_core_type)
        if ocr_core_type == "PaddleOCR":
            return ocr_core.runBytes(file_bytes)
        else:
            return ocr_core.runBytes(file_bytes)
    finally:
        if ocr_core:
            release_ocr_instance(ocr_core_type, ocr_core)


def process_image_file_task(image_path, ocr_core_type: str = "RapidOCR"):
    """优化版文件图片处理"""
    ocr_core = None
    try:
        ocr_core = get_ocr_instance(ocr_core_type)
        if ocr_core_type == "PaddleOCR":
            return ocr_core.run(image_path)
        else:
            return ocr_core.run(image_path)
    finally:
        if ocr_core:
            release_ocr_instance(ocr_core_type, ocr_core)

def pdf_to_images(pdf_path, output_dir, file_name):
    """
    pdf转图片
    :param pdf_path: PDF文件路径
    :param output_dir: 输出目录
    :param file_name: 文件名前缀
    :return: 图片文件路径列表
    """
    import fitz
    pdf_doc = fitz.open(pdf_path)
    file_names = []
    for page_index in range(pdf_doc.page_count):
        page = pdf_doc[page_index]
        zoom = (1.3, 1.3)
        mat = fitz.Matrix(zoom[0], zoom[1])
        pix = page.get_pixmap(matrix=mat, alpha=False)
        temp_file_name = file_name + str(page_index) + '.png'
        image_abs_path = os.path.join(output_dir, temp_file_name)
        pix.save(image_abs_path)
        file_names.append(image_abs_path)
    return file_names


def get_file_extension(filename):
    """
    获取文件扩展名
    :param filename:
    :return:
    """
    return os.path.splitext(filename)[1].lower()


def allowed_file(filename):
    """
    判断文件是否为允许的类型
    :param filename:
    :return:
    """
    return get_file_extension(filename) in ALLOWED_EXTENSIONS

@post("/ocr_optimized")
def ocr_optimized():
    """OCR接口 适合100-200张图片"""
    upload_files = request.files.getall('file')
    if not upload_files:
        return json.dumps({"code": 400, "msg": "请上传文件"})

    start_time = time()
    tasks = []

    with tempfile.TemporaryDirectory() as tmp_dir:
        # 任务准备阶段
        for upload_file in upload_files:
            if not allowed_file(upload_file.raw_filename):
                tasks.append(OCRTask(
                    type='error',
                    filename=upload_file.raw_filename,
                    error=f"不支持的文件类型: {get_file_extension(upload_file.raw_filename)}"
                ))
                continue

            file_extension = get_file_extension(upload_file.raw_filename)
            ocr_type = request.query.get("ocr_type", "RapidOCR")

            if file_extension in IMAGE_EXTENSIONS:
                tasks.append(OCRTask(
                    type='image_memory',
                    filename=upload_file.raw_filename,
                    data=upload_file.file.read(),
                    ocr_type=ocr_type
                ))
            elif file_extension == '.pdf':
                file_path = os.path.join(tmp_dir, upload_file.raw_filename)
                upload_file.save(file_path)
                image_files = pdf_to_images(file_path, tmp_dir,
                                                      upload_file.raw_filename[:-4])
                for i, image_path in enumerate(image_files):
                    tasks.append(OCRTask(
                        type='image_file',
                        filename=f"{upload_file.raw_filename}_page_{i + 1}",
                        data=image_path,
                        ocr_type=ocr_type,
                        original_filename=upload_file.raw_filename,
                        page_num=i + 1
                    ))

        # 并行处理阶段
        results = []
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            futures = []
            for task in tasks:
                if task.type == 'image_memory':
                    futures.append(executor.submit(
                        partial(process_image_memory_task, ocr_core_type=task.ocr_type),
                        task.data
                    ))
                elif task.type == 'image_file':
                    futures.append(executor.submit(
                        partial(process_image_file_task, ocr_core_type=task.ocr_type),
                        task.data
                    ))

            for future in as_completed(futures):
                try:
                    result = future.result()
                    results.append(result)
                except Exception as e:
                    results.append({"error": str(e)})

    total_time = time() - start_time
    print(f"处理完成，总耗时: {total_time:.2f}秒")
    return json.dumps({
        "code": 200,
        "msg": "success",
        "data": results,
        "performance": {
            "total_time": total_time,
            "avg_time_per_task": total_time / len(tasks) if tasks else 0,
            "workers_used": MAX_WORKERS
        }
    })

# @route('/ocr_test')
# def ocr_test():
#     """
#     测试ocr 是否可用
#     :return:
#     """
#     try:
#         # 支持的ocr
#         supported_ocr = []
#         for ocr_core_type in ["PaddleOCR", "RapidOCR"]:
#             result = process_image_file_task(TEST_OCR_FILE_PATH,ocr_core_type=ocr_core_type)
#             print(result)
#             if result:
#                 supported_ocr.append(ocr_core_type)
#         return json.dumps({"code": 200, "msg": "ocr 可用", "data": supported_ocr})
#     except Exception as e:
#         return json.dumps({"code": 500, "msg": f"ocr 不可用: {str(e)}"})
# @route('/ocr_test')
# def ocr_test():
#     """
#     测试ocr 是否可用
#     :return:
#     """
#     try:
#         # 支持的ocr
#         supported_ocr = []
#         for ocr_core_type in ["PaddleOCR", "RapidOCR"]:
#             result = process_image_file_task(TEST_OCR_FILE_PATH,ocr_core_type=ocr_core_type)
#             print(result)
#             if result:
#                 supported_ocr.append(ocr_core_type)
#         return json.dumps({"code": 200, "msg": "ocr 可用", "data": supported_ocr})
#     except Exception as e:
#         return json.dumps({"code": 500, "msg": f"ocr 不可用: {str(e)}"})

@post("/ocr_sse_optimized")
def ocr_sse_optimized():
    """OCR SSE接口 文件过大时使用 后期会接受一个本地路径 直接读取"""
    response.content_type = 'text/event-stream'
    response.cache_control = 'no-cache'

    def event_stream():
        upload_files = request.files.getall('file')
        if not upload_files:
            yield f"data: {json.dumps({'code': 400, 'msg': '请上传文件'})}\n\n"
            yield "event: end\ndata: {}\n\n"
            return

        start_time = time()
        tasks = []

        with tempfile.TemporaryDirectory() as tmp_dir:
            # 任务准备阶段
            for upload_file in upload_files:
                if not allowed_file(upload_file.raw_filename):
                    yield f"data: {json.dumps({'filename': upload_file.raw_filename, 'status': 'error', 'error': f'不支持的文件类型: {get_file_extension(upload_file.raw_filename)}'})}\n\n"
                    continue

                file_extension = get_file_extension(upload_file.raw_filename)
                ocr_type = request.query.get("ocr_type", "RapidOCR")

                if file_extension in IMAGE_EXTENSIONS:
                    tasks.append(OCRTask(
                        type='image_memory',
                        filename=upload_file.raw_filename,
                        data=upload_file.file.read(),
                        ocr_type=ocr_type
                    ))
                elif file_extension == '.pdf':
                    yield f"data: {json.dumps({'filename': upload_file.raw_filename, 'status': 'processing', 'message': '正在转换PDF为图片'})}\n\n"
                    file_path = os.path.join(tmp_dir, upload_file.raw_filename)
                    upload_file.save(file_path)
                    image_files = pdf_to_images(file_path, tmp_dir, upload_file.raw_filename[:-4])
                    for i, image_path in enumerate(image_files):
                        tasks.append(OCRTask(
                            type='image_file',
                            filename=f"{upload_file.raw_filename}_page_{i + 1}",
                            data=image_path,
                            ocr_type=ocr_type,
                            original_filename=upload_file.raw_filename,
                            page_num=i + 1
                        ))

            yield f"data: {json.dumps({'status': 'processing', 'total_tasks': len(tasks), 'workers': MAX_WORKERS})}\n\n"

            # 并行处理阶段
            completed = 0
            with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
                futures_to_task = {}
                for task in tasks:
                    if task.type == 'image_memory':
                        futures = executor.submit(
                            partial(process_image_memory_task, ocr_core_type=task.ocr_type),
                            task.data
                        )
                    elif task.type == 'image_file':
                        futures = executor.submit(
                            partial(process_image_file_task, ocr_core_type=task.ocr_type),
                            task.data
                        )
                    futures_to_task[futures] = task

                for future in as_completed(futures_to_task):
                    task = futures_to_task[future]
                    try:
                        result = future.result()
                        completed += 1
                        if isinstance(result, dict):
                            result.update({
                                'filename': task.filename,
                                'status': 'completed',
                                'progress': f'{completed}/{len(tasks)}'
                            })
                        if task.type == 'image_memory':
                            yield f"data: {json.dumps(result)}\n\n"
                        else:  # image_file (PDF页面)
                            result.update({ 'original_filename': task.original_filename, 'page_num': task.page_num, })
                            yield f"data: {json.dumps(result)}\n\n"
                    except Exception as e:
                        completed += 1
                        if task.type == 'image_memory':
                            yield f"data: {json.dumps({'filename': task.filename, 'error': str(e), 'completed': completed, 'total': len(tasks)})}\n\n"
                        else:  # image_file
                            yield f"data: {json.dumps({'filename': task.filename, 'original_filename': task.original_filename, 'page_num': task.page_num, 'error': str(e), 'completed': completed, 'total': len(tasks)})}\n\n"

            yield f"data: {json.dumps({'status': 'all_completed', 'total_time': time() - start_time})}\n\n"
            yield "event: end\ndata: {}\n\n"

    return event_stream()



def start_bottle(port: int = 9911, memory_size=300 * 1024 * 1024, hot_start_bottle=True):
    """
    bottle 启动函数
    :param port:
    :param memory_size: 300MB
    :param hot_start_bottle: 是否预启动OCR进程池
    :return:
    """
    # 设置文件大小限制
    BaseRequest.MEMFILE_MAX = memory_size  # 300MB
    # 预启动

    if hot_start_bottle:
        init_ocr_pool()
    run(port=port)

if __name__ == '__main__':
    start_bottle()